# Auto-generated unit tests for trt_utils

import sys
import types
import importlib
import builtins
import onnx
import onnx.helper as oh
import onnx_graphsurgeon as gs
import pytest

# Detect source module path
# Prefer modelopt.onnx.trt_utils; fallback to importing via tests path if needed.
def _import_trt_utils_with_stubbed_tensorrt(stub=None):
    if stub is None:
        # Minimal TensorRT stub with DataType enums required by infer_types_shapes
        stub = types.ModuleType("tensorrt")
        class _DT:
            pass
        # Use distinct singleton objects for mapping keys
        for name in ["float32","float16","bfloat16","int4","int8","uint8","int32","int64","bool","fp8","fp4"]:
            setattr(stub, name, object())
        stub.DataType = _DT
        # Minimal Logger enum placeholders to avoid accidental imports if functions call them
        class _Logger:
            WARNING = 2
        stub.Logger = _Logger
        class _NDCF:
            STRONGLY_TYPED = 0
        stub.NetworkDefinitionCreationFlag = _NDCF
    sys.modules["tensorrt"] = stub
    # Fresh import to bind the stubbed tensorrt
    try:
        mod = importlib.import_module("modelopt.onnx.trt_utils")
    except ModuleNotFoundError:
        # Fallback: try relative import if project layout differs
        mod = importlib.import_module("trt_utils")
    return mod

@pytest.fixture
def trt_utils_module(monkeypatch):
    # Ensure a clean tensorrt stub per test
    if "tensorrt" in sys.modules:
        del sys.modules["tensorrt"]
    return _import_trt_utils_with_stubbed_tensorrt()

def _simple_graph_with_custom_op(op_type="MyCustomOp", num_inps=2, num_outs=1):
    """
    Build a tiny ONNX model using GraphSurgeon where a node with op_type exists.
    Inputs/outputs are tensors without specific types/shapes; sufficient for traversal.
    """
    inputs = [gs.Variable(f"in{i}", dtype=None, shape=None) for i in range(num_inps)]
    outputs = [gs.Variable(f"out{i}", dtype=None, shape=None) for i in range(num_outs)]
    node = gs.Node(op=op_type, inputs=inputs, outputs=outputs, name=f"{op_type}_node")
    graph = gs.Graph(nodes=[node], inputs=inputs, outputs=outputs)
    model = gs.export_onnx(graph)
    return model, node

class TestInterpretTrtPluginsPrecisionFlag:
    def test_single_precision_fp16_casts_all_io(self, trt_utils_module):
        # Model has one custom op with 2 inputs and 3 outputs
        model, _ = _simple_graph_with_custom_op(op_type="MyOp", num_inps=2, num_outs=3)
        cast, quant = trt_utils_module.interpret_trt_plugins_precision_flag(
            model, ["MyOp:fp16"], quantize_mode="int8"
        )
        assert cast == {"MyOp": {"inp": [0,1], "out": [0,1,2]}}
        assert quant == {}

    def test_single_precision_int8_quantizes_all_io_respects_quantize_mode(self, trt_utils_module):
        # Two nodes of same type but different arities; function uses max across nodes
        model1, _ = _simple_graph_with_custom_op("MyQ", 1, 1)
        model2, _ = _simple_graph_with_custom_op("MyQ", 3, 2)
        # Merge graphs by concatenating nodes in a new graph
        g = gs.import_onnx(model1)
        g2 = gs.import_onnx(model2)
        graph = gs.Graph(nodes=g.nodes + g2.nodes, inputs=g.inputs + g2.inputs, outputs=g.outputs + g2.outputs)
        merged = gs.export_onnx(graph)

        cast, quant = trt_utils_module.interpret_trt_plugins_precision_flag(
            merged, ["MyQ:fp8"], quantize_mode="int8"
        )
        # Even though fp8 requested, it should align to quantize_mode=int8 and select max arity: 3 inps, 2 outs
        assert "MyQ" in quant and quant["MyQ"]["inp"] == [0,1,2] and quant["MyQ"]["out"] == [0,1]
        assert cast == {}

    def test_explicit_io_precisions_cast_and_quantize(self, trt_utils_module):
        # 2 inputs, 3 outputs
        model, _ = _simple_graph_with_custom_op("OpX", 2, 3)
        cast, quant = trt_utils_module.interpret_trt_plugins_precision_flag(
            model, ["OpX:[fp16,fp32]:[fp16,int8,fp32]"], quantize_mode="int8"
        )
        # Cast: inputs with fp16 -> [0]; outputs with fp16 or fp32 -> [0,2]
        assert cast["OpX"]["inp"] == [0]
        assert cast["OpX"]["out"] == [0,2]
        # Quantize: outputs with int8 -> [1]
        assert quant["OpX"]["inp"] == []
        assert quant["OpX"]["out"] == [1]

    def test_mismatched_counts_raise(self, trt_utils_module):
        model, _ = _simple_graph_with_custom_op("OpY", 2, 2)
        with pytest.raises(AssertionError):
            trt_utils_module.interpret_trt_plugins_precision_flag(
                model, ["OpY:[fp16]:[fp16,fp16]"], quantize_mode="int8"
            )
        with pytest.raises(AssertionError):
            trt_utils_module.interpret_trt_plugins_precision_flag(
                model, ["OpY:[fp16,fp16]:[fp16]"], quantize_mode="int8"
            )

    def test_unknown_op_type_is_skipped(self, trt_utils_module, caplog):
        model, _ = _simple_graph_with_custom_op("Existing", 1, 1)
        cast, quant = trt_utils_module.interpret_trt_plugins_precision_flag(
            model, ["MissingOp:fp16"], quantize_mode="int8"
        )
        assert cast == {}
        assert quant == {}

class TestSetTrtPluginDomain:
    def test_sets_domain_and_opset_once(self, trt_utils_module):
        # Graph with custom and vanilla nodes
        inputs = [gs.Variable("a"), gs.Variable("b")]
        out1 = gs.Variable("c")
        out2 = gs.Variable("d")
        custom = gs.Node(op="CustomTRT", inputs=[inputs[0]], outputs=[out1], name="custom1")
        add = gs.Node(op="Add", inputs=[out1, inputs[1]], outputs=[out2], name="add1")
        graph = gs.Graph(nodes=[custom, add], inputs=inputs, outputs=[out2])
        model = gs.export_onnx(graph)

        updated = trt_utils_module.set_trt_plugin_domain(model, custom_ops={"CustomTRT"})
        g = gs.import_onnx(updated)
        # Domain should be set on custom op only
        custom_nodes = [n for n in g.nodes if n.op == "CustomTRT"]
        assert len(custom_nodes) == 1 and custom_nodes[0].domain == "trt.plugins"
        add_nodes = [n for n in g.nodes if n.op == "Add"]
        assert add_nodes and (add_nodes[0].domain is None or add_nodes[0].domain == "")

        # TRT domain opset should be appended
        domains = {(opset.domain, opset.version) for opset in updated.opset_import}
        assert ("trt.plugins", 1) in domains

class TestInferTypesShapes:
    def test_sets_elem_type_and_shape_on_existing_value_info(self, trt_utils_module):
        # Create a simple model: X + Y -> Z; add Z to value_info without type/shape
        X = oh.make_tensor_value_info("X", onnx.TensorProto.FLOAT, [1, 3])
        Y = oh.make_tensor_value_info("Y", onnx.TensorProto.FLOAT, [1, 3])
        Z = oh.make_tensor_value_info("Z", 0, [])  # elem_type unset (0), shape empty
        node = oh.make_node("Add", inputs=["X", "Y"], outputs=["Z"], name="add_node")
        graph = oh.make_graph([node], "g", [X, Y], [oh.make_tensor_value_info("O", onnx.TensorProto.FLOAT, [1,3])])
        # Attach Z as value_info (not output)
        graph.value_info.extend([Z])
        model = oh.make_model(graph)

        # all_tensor_info provides dtype (TRT) and shape (with dynamic str dim)
        trt = sys.modules["tensorrt"]
        info = {"Z": {"dtype": getattr(trt, "float16"), "shape": [1, "unk", 3]}}

        updated = trt_utils_module.infer_types_shapes(model, info)

        # Fetch updated Z value_info
        def _get_vi(m, name):
            for vi in m.graph.value_info:
                if vi.name == name:
                    return vi
            return None

        z_vi = _get_vi(updated, "Z")
        assert z_vi is not None
        assert z_vi.type.tensor_type.elem_type == onnx.TensorProto.FLOAT16
        dims = z_vi.type.tensor_type.shape.dim
        assert [d.dim_value if d.HasField("dim_value") else d.dim_param for d in dims] == [1, "unk", 3]

    def test_appends_value_info_when_missing(self, trt_utils_module):
        # Node output W not present in value_info; expect function to append it.
        A = oh.make_tensor_value_info("A", onnx.TensorProto.FLOAT, [1, 2])
        B = oh.make_tensor_value_info("B", onnx.TensorProto.FLOAT, [1, 2])
        node = oh.make_node("Mul", inputs=["A", "B"], outputs=["W"], name="mul_node")
        graph = oh.make_graph([node], "g2", [A, B], [])
        model = oh.make_model(graph)

        trt = sys.modules["tensorrt"]
        info = {"W": {"dtype": getattr(trt, "int8"), "shape": ["unk", 2]}}

        updated = trt_utils_module.infer_types_shapes(model, info)

        names = [vi.name for vi in updated.graph.value_info]
        assert "W" in names
        # Ensure elem type is INT8
        vi = next(vi for vi in updated.graph.value_info if vi.name == "W")
        assert vi.type.tensor_type.elem_type == onnx.TensorProto.INT8
        dims = vi.type.tensor_type.shape.dim
        assert [d.dim_value if d.HasField("dim_value") else d.dim_param for d in dims] == ["unk", 2]